package cn.mybatis.plugin.pagehelper;

import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;
import java.lang.reflect.Method;
import java.sql.SQLException;
import java.util.Map;

/**
 * @author: 张澎_9970
 * @date: 2022.1.17 下午 11:20
 */
@Intercepts({@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})})
public class PageInterceptor implements Interceptor {

    private JdbcTemplate jdbcTemplate;

    public PageInterceptor(DataSource dataSource) {
        jdbcTemplate = new JdbcTemplate(dataSource);
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Page<Object> page = PageHelper.getPage();
        if (page == null) {
            return invocation.proceed();
        }

        Executor executor = (Executor) invocation.getTarget();
        Method method = invocation.getMethod();
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object parameters = args[1];
        RowBounds bounds = (RowBounds) args[2];

        ResultHandler resultHandler = (ResultHandler) args[3];
        try {
            if (args.length == 6) {
                CacheKey cacheKey = (CacheKey) args[4];
                BoundSql boundSql = (BoundSql) args[5];

                // 总数查询
                Long count = queryCount(executor, ms, parameters, bounds, resultHandler, boundSql);
                page.setTotal(count.intValue());

                // 执行分页查询
                final String pageSql = PageSqlDialectHolder.getPageSqlDialect().getPageSql(boundSql.getSql(), page);
                final BoundSql pageBoundSql = new BoundSql(ms.getConfiguration(), pageSql, boundSql.getParameterMappings(), parameters);
                page.addAll(executor.query(ms, parameters, bounds, resultHandler, cacheKey, pageBoundSql));
                return page;

            } else {// 4个参数
                // TODO
                System.out.println("4个参数");
            }
        } finally {
            PageHelper.clearPage();
        }

        return invocation.proceed();
    }

    private Long queryCount(Executor executor, MappedStatement ms, Object parameters, RowBounds bounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        String countSql = PageSqlDialectHolder.getPageSqlDialect().getCountSql(boundSql.getSql());

        // // TODO: 2022.1.18
        Map<String, Object> map = jdbcTemplate.queryForMap(countSql);
        return (Long) map.values().stream().findFirst().orElse(0);
    }

}
